from sentence_transformers import SentenceTransformer

embedding_dim =512
# 加载BGE模型
embedding = SentenceTransformer("E:/huaqing/ai_25021/项目/智能客服工单分类系统/bge_small_zh/bge_small_zh")  # 绝对路径

# 冻结参数
for p in embedding.parameters():
    # 关闭自动求梯度追踪
    p.requires_grad = False

# 分词编号函数
# texts: 多句话的文本列表[“你好”,"你是谁？","你能帮我做什么呢？"]
# max_len: 其实就是在大模型中见到的max_tokens限制，限制输入输出最大词元数量
def tokenizer(texts, max_len=50):
    # 分词
    result = embedding.tokenizer(texts,                     # 多句话
                                 add_special_tokens=False,  # 不添加新词
                                 max_length=max_len,        # 最大长度
                                 padding="max_length",      # 达不到最大长度，填充
                                 truncation=True,           # 超过最大长度，截断
                                 return_tensors="pt",       # 返回pytorch张量
                                 )
    ids = result["input_ids"]
    padding_mask = result["attention_mask"].float()
    padding_mask[padding_mask == 0] = float('-inf')
    padding_mask[padding_mask == 1] = 0.0
    return ids, padding_mask

if __name__ == '__main__':
    texts = "1+1=2"
    print(tokenizer(texts, max_len=10))